package algorithm.graph.adjacencymatrix;

import algorithm.graph.base.Edge;
import algorithm.graph.base.GraphUtil;

import java.util.*;
import java.util.stream.Collectors;
import java.util.stream.IntStream;

public class Prim {
    AdjacencyMatrixGraph graph;

    public Prim(AdjacencyMatrixGraph graph) {
        this.graph = graph;
    }

    public AdjacencyMatrixGraph buildMinEdgeGraph() {
        List<Edge> edgeList = buildMinEdgeGraphEdgeList(graph);
        AdjacencyMatrixGraph result = new AdjacencyMatrixGraph(graph.getNodeCount());
        result.initNonDirectedGraphEdges(edgeList);
        return result;
    }

    private static List<Edge> buildMinEdgeGraphEdgeList(AdjacencyMatrixGraph graph) {
        List<Edge> unvisitedEdgeList = GraphUtil.getNonDirectedGraphEdgeList(graph);
        List<Edge> validEdgeList = new ArrayList<>();
        List<Integer> unvisitedNodeList = IntStream.range(0, graph.getNodeCount()).boxed().collect(Collectors.toCollection(ArrayList::new));
        List<Integer> visitedNodeList = new ArrayList<>();
        while (!unvisitedNodeList.isEmpty()) {
            if (visitedNodeList.size() == 0) {
                visitedNodeList.add(unvisitedNodeList.remove(0));
            } else {
                Edge nextEdge = findMinWeightEdge(visitedNodeList, unvisitedNodeList, unvisitedEdgeList);
                if (nextEdge != null) {
                    visitedNodeList.add(nextEdge.to());
                    GraphUtil.removeItemFromList(unvisitedNodeList, it -> nextEdge.to() == it);
                    updateInvalidEdgeList(unvisitedEdgeList, visitedNodeList, nextEdge);
                    validEdgeList.add(nextEdge);
                } else {
                    break;
                }
            }
        }
        if (!unvisitedNodeList.isEmpty()) {
            throw new RuntimeException("invalid graph");
        }
        return validEdgeList;
    }

    private static void updateInvalidEdgeList(List<Edge> unvisitedEdgeList, List<Integer> visitedNodeList, Edge nextEdge) {
        GraphUtil.removeItemFromList(unvisitedEdgeList, (it) -> {
            return visitedNodeList.contains(it.from()) && nextEdge.to() == it.to();
        });
    }

    private static Edge findMinWeightEdge(List<Integer> visitedNodeList, List<Integer> unvisitedNodeList, List<Edge> unvisitedEdgeList) {
        Edge minWeightEdge = null;
        for (Edge edge : unvisitedEdgeList) {
            if (visitedNodeList.contains(edge.from()) && unvisitedNodeList.contains(edge.to())) {
                if (minWeightEdge == null) {
                    minWeightEdge = edge;
                } else {
                    if (edge.weight() < minWeightEdge.weight()) {
                        minWeightEdge = edge;
                    }
                }
            }
        }
        return minWeightEdge;
    }
}
